package com.huawei.admins.platform.gateway.core.component.router;

import com.huawei.admins.platform.gateway.core.component.cache.GatewayLocalCache;
import com.huawei.admins.platform.gateway.core.config.GlobalCorsProperties;
import jakarta.validation.constraints.NotNull;
import org.springframework.core.env.Environment;
import org.springframework.web.reactive.handler.AbstractHandlerMapping;
import org.springframework.web.server.ServerWebExchange;
import reactor.core.publisher.Mono;

import java.util.function.Function;

import static com.huawei.admins.platform.gateway.core.util.ServerWebExchangeUtils.GATEWAY_PREDICATE_ROUTE_ATTR;
import static com.huawei.admins.platform.gateway.core.util.ServerWebExchangeUtils.GATEWAY_REACTOR_CONTEXT_ATTR;
import static com.huawei.admins.platform.gateway.core.util.ServerWebExchangeUtils.GATEWAY_ROUTE_ATTR;

public class RoutePredicateHandlerMapping extends AbstractHandlerMapping {

    private final GatewayLocalCache<Router> routeGatewayLocalCache = new GatewayLocalCache<>();

    private final RouterLocator routerLocator;

    public RoutePredicateHandlerMapping(GlobalCorsProperties globalCorsProperties, Environment environment, RouterLocator routerLocator) {
        this.routerLocator = routerLocator;
        setOrder(environment.getProperty("spring.cloud.gateway.handler-mapping.order", Integer.class, 1));
        setCorsConfigurations(globalCorsProperties.getCorsConfigurations());
    }

    @NotNull
    @Override
    protected Mono<?> getHandlerInternal(ServerWebExchange exchange) {
        return Mono.deferContextual(contextView -> {
            exchange.getAttributes().put(GATEWAY_REACTOR_CONTEXT_ATTR, contextView);
            return lookupRoute(exchange)
                    .map((Function<Router, ?>) r -> {
                        exchange.getAttributes().remove(GATEWAY_PREDICATE_ROUTE_ATTR);
                        if (logger.isDebugEnabled()) {
                            logger.debug("Mapping [" + getExchangeDesc(exchange) + "] to " + r);
                        }

                        exchange.getAttributes().put(GATEWAY_ROUTE_ATTR, r);
                        return r;
                    }).switchIfEmpty(Mono.empty().then(Mono.fromRunnable(() -> {
                        exchange.getAttributes().remove(GATEWAY_PREDICATE_ROUTE_ATTR);
                        if (logger.isTraceEnabled()) {
                            logger.trace("No RouteDefinition found for [" + getExchangeDesc(exchange) + "]");
                        }
                    })));
        });
    }

    protected Mono<Router> lookupRoute(ServerWebExchange exchange) {

        String rawPath = exchange.getRequest().getPath().value();
        String routePath = normalizeRoutePath(rawPath);

        Router route = routeGatewayLocalCache.get(routePath);
        if (route != null) {
            return Mono.just(route);
        }
        return lookupRouteInternal(exchange).doOnSuccess(res -> {
            if (res != null) {
                routeGatewayLocalCache.set(routePath, res);
            }
        });
    }

    protected Mono<Router> lookupRouteInternal(ServerWebExchange exchange) {
        return this.routerLocator.getRouters()
                .concatMap(route -> Mono.just(route).filterWhen(r -> {
                            // add the current route we are testing
                            exchange.getAttributes().put(GATEWAY_PREDICATE_ROUTE_ATTR, r.getId());
                            return r.getPredicate().apply(exchange);
                        })
                        .doOnError(e -> logger.error("Error applying predicate for route: " + route.getId(), e))
                        .onErrorResume(e -> Mono.empty()))
                .next()
                .map(route -> {
                    if (logger.isDebugEnabled()) {
                        logger.debug("Route matched: " + route.getId());
                    }
                    return route;
                });
    }

    private String normalizeRoutePath(String rawPath) {
        return rawPath.substring(0, rawPath.lastIndexOf("/"));
    }

    private String getExchangeDesc(ServerWebExchange exchange) {
        return "Exchange: " +
                exchange.getRequest().getMethod() +
                " " +
                exchange.getRequest().getURI();
    }
}
